// Fill out your copyright notice in the Description page of Project Settings.
#include "HightScreenshotLib.h"
#include "Engine/Canvas.h"
#include "ScreenshotViewport.h"
#include "Misc/FileHelper.h"
#include "IImageWrapper.h"
#include "IImageWrapperModule.h"

#include "Modules/ModuleManager.h"
//#include "Engine/Public/EngineModule.h"
#include "ScreenshotViewportClient.h"
#include "Engine/World.h"
#include "Engine/Engine.h"
#include "EngineUtils.h"
#include "Engine/PostProcessVolume.h"

#include "PreviewScene.h"
#include "ImageUtils.h"

//管理世界和ViewportClient
//保证每个worl只创建一个ViewportClient
//在世界被销毁时自动销毁创建的ViewportClient
class FHoldViewport 
{
public:
	FHoldViewport(FHoldViewport&) = delete;
	FHoldViewport& operator=(const FHoldViewport&) = delete;

	static FScreenshotViewportClient* GetViewportClient(AActor* targetActo);

private:

	TMap<uint32, FScreenshotViewportClient*> worldID_ViewportClient;

	FHoldViewport();

	void  OnWorldCleanup(UWorld* world, bool bSessionEnded, bool bCleanupResources);


};

FHoldViewport::FHoldViewport() 
{
	FWorldDelegates::OnPostWorldCleanup.AddRaw(this,&FHoldViewport::OnWorldCleanup);
}

FScreenshotViewportClient* FHoldViewport::GetViewportClient(AActor* targetActo)
{
	static FHoldViewport holdWorl;
	check(targetActo);
	UWorld* wrold = targetActo->GetWorld();
	uint32 worldID = wrold->GetUniqueID();
	if (holdWorl.worldID_ViewportClient.Contains(worldID))
	{
		return holdWorl.worldID_ViewportClient[worldID];
	}
	else
	{
		FScreenshotViewportClient* customViewportClient = new FScreenshotViewportClient();
		customViewportClient->InitViewportClient(wrold);
		holdWorl.worldID_ViewportClient.Add(worldID, customViewportClient);
		return customViewportClient;
	}
}

void FHoldViewport::OnWorldCleanup(UWorld* world, bool bSessionEnded, bool bCleanupResources)
{
	uint32 worldID = world->GetUniqueID();
	if (worldID_ViewportClient.Contains(worldID))
	{
		FScreenshotViewportClient* customViewportClient = worldID_ViewportClient[worldID];
		delete customViewportClient;
		worldID_ViewportClient.Remove(worldID);
		customViewportClient = nullptr;
	}
}



bool UHightScreenshotLib::SimpleScreenshot(
	AActor* targetActo, FVector point, FRotator rotator
	, int sizeX, int sizeY, TArray<FColor>& outColor, float fov, int inFrameDelay
	, EScreenshotShowFlags flags)
{

	outColor.Reset();

	if (!targetActo || sizeX * sizeY <= 0)return false;


	FScreenshotInfo info;
	info.worldTargetActor = targetActo;
	info.FrameDelay = inFrameDelay;
	info.flags = flags;
	info.outImageInfo.format = EOutImageFormat::JPEG;
	info.outImageInfo.width = sizeX;
	info.outImageInfo.height = sizeY;

	FMinimalViewInfo  ViewInfo  = FMinimalViewInfo();
	ViewInfo.FOV = info.minimalViewInfo.DesiredFOV = FMath::Max(16.0f, fov);
	ViewInfo.DesiredFOV = ViewInfo.FOV;
	ViewInfo.Location = point;
	ViewInfo.Rotation = rotator;
	ViewInfo.bUseFieldOfViewForLOD = false;
	ViewInfo.AspectRatio = float(sizeX) / float(sizeY);

	info.minimalViewInfo = ViewInfo;

	GetFirstPostProcessSettings(targetActo
		, info.minimalViewInfo.PostProcessSettings
		, info.minimalViewInfo.PostProcessBlendWeight);


	return Screenshot(info,outColor);
}



bool  UHightScreenshotLib::Screenshot(FScreenshotInfo info, TArray<FColor>& outColor)
{
	UWorld* ScreenshotWorld = info.worldTargetActor->GetWorld();
	int sizeX = info.outImageInfo.width;
	int sizeY = info.outImageInfo.height;


	UGameViewportClient* gameViewportClient = ScreenshotWorld->GetGameViewport();

	FScreenshotViewportClient* customViewportClient = FHoldViewport::GetViewportClient(info.worldTargetActor);


	customViewportClient->SetEngineShowFlags(*gameViewportClient->GetEngineShowFlags());



	//ViewportClient* drawView = gameViewportClient;
	FViewportClient* drawView =  customViewportClient;



	FStoneDummyViewport* DummyViewport = new FStoneDummyViewport(drawView);
	DummyViewport->SetTextureSize(sizeX, sizeY);

	customViewportClient->InitRenderTask(info, DummyViewport);


	BeginInitResource((FRenderResource*)DummyViewport);

	uint32 FrameDelay = FMath::Max(2, info.FrameDelay);

	if (info.flags != EScreenshotShowFlags::Lit)
	{
		//其他通道不需要叠加太多
		FrameDelay = 2;
	}

	FStoneDummyViewport::EndFrame(nullptr);

	while (FrameDelay)
	{
		FStoneDummyViewport::BeginFrame(DummyViewport);

		FCanvas Canvas(DummyViewport, NULL, ScreenshotWorld, ScreenshotWorld->GetFeatureLevel());
		{
			drawView->Draw(DummyViewport, &Canvas);
		}
		Canvas.Flush_GameThread();

		FStoneDummyViewport::EndFrame(DummyViewport);

		FlushRenderingCommands();

		--FrameDelay;
	}


	bool bScreenshotSuccessful = GetViewportScreenShot(DummyViewport, outColor, FIntRect(0, 0, sizeX, sizeY));

	BeginReleaseResource((FRenderResource*)DummyViewport);
	FlushRenderingCommands();
	delete DummyViewport;
	FStoneDummyViewport::BeginFrame(nullptr);

	customViewportClient->FlushRenderingRenderTask();

	//强行设置A通道  todo 透明度信息未能读取到
	int num = outColor.Num();
	for (size_t i = 0; i < num; i++)
	{
		outColor[i].A = 255;
	}
	return bScreenshotSuccessful;
}

bool  UHightScreenshotLib::CompressTexture(int sizeX, int sizeY, const TArray<FColor>& Colors
	, EOutImageFormat format, TArray<uint8>& outData, int q)
{
	   
	if (sizeX*sizeY != Colors.Num() || Colors.Num() <= 1) 
	{
		return false;
	}

	IImageWrapperModule& ImageWrapperModule = FModuleManager::LoadModuleChecked<IImageWrapperModule>(FName("ImageWrapper"));
	EImageFormat imageformat = EImageFormat::JPEG;
	switch (format)
	{	

	case EOutImageFormat::JPEG:
		imageformat = EImageFormat::JPEG;
		break;
	case EOutImageFormat::PNG:
		imageformat = EImageFormat::PNG;
		break;
	case EOutImageFormat::BMP:
		imageformat = EImageFormat::BMP;
		break;
	case EOutImageFormat::WepP:
		imageformat = EImageFormat::JPEG;
	default:
		imageformat = EImageFormat::JPEG;
		break;
	}

	TSharedPtr<IImageWrapper> ImageWrapper = ImageWrapperModule.CreateImageWrapper(imageformat);
	if (!ImageWrapper.IsValid()) return false;

	int TextureBufferMemorySize = sizeX * sizeY * sizeof(FColor);

	bool b = ImageWrapper->SetRaw(Colors.GetData()
		, TextureBufferMemorySize
		, sizeX, sizeY
		, ERGBFormat::BGRA, 8);

	outData = TArray<uint8>(ImageWrapper->GetCompressed());

	return b;
}

bool UHightScreenshotLib::SaveToDisk(FString path, const TArray<uint8>& data)
{
	if (path.Len()<=8|| data.Num() <= 1)
	{
		return false;
	}
	return FFileHelper::SaveArrayToFile(data, *path);
}

void  UHightScreenshotLib::GetFirstPostProcessSettings(UObject* WorldContextObject
	, FPostProcessSettings& outPostProcessSettings,float& PostProcessBlendWeight)
{
	if (GEngine&&WorldContextObject) 
	{
		UWorld* World = GEngine->GetWorldFromContextObject(WorldContextObject, EGetWorldErrorMode::LogAndReturnNull);
		if (World)
		{
			{
				for (TActorIterator<APostProcessVolume> It(World, APostProcessVolume::StaticClass()); It; ++It)
				{
					APostProcessVolume* Actor = *It;
					if (IsValid(Actor))
					{
						outPostProcessSettings = Actor->Settings;
						PostProcessBlendWeight = Actor->BlendWeight;
						break;
					}
				}
			}
		}
	}
}

bool UHightScreenshotLib::MirroTexture(int sizeX, int sizeY
	, const TArray<FColor>& inColors, TArray<FColor>& outColors, bool mirrorX, bool mirrorY)
{
	int len = inColors.Num();
	outColors = inColors;
	if (sizeX* sizeY!= len)
	{
		return true;
	}
	if (!mirrorX&&!mirrorY)
	{
		return true;
	}
	FColor cb;
	int iy,a,b;
	if (mirrorY)
	{
		int hy = sizeY / 2;
		for (size_t y = 0; y < hy; y++)
		{
			iy = sizeY - y-1;
			for (size_t x = 0; x < sizeX; x++)
			{
				a = iy * sizeX + x;
				b = y * sizeX + x;
				cb = outColors[b];
				outColors[b] = outColors[a];
				outColors[a] = cb;
			}
		}
	}
	if (mirrorX)
	{
		int hx = sizeX / 2;
		for (size_t y = 0; y < sizeY ; y++)
		{
			for (size_t x = 0; x < hx; x++)
			{
				a = y * sizeX + x;
				b = y * sizeX + (sizeX - x - 1);
				cb = outColors[b];
				outColors[b] = outColors[a];
				outColors[a] = cb;
			}
		}
	}
	return true;
}
UTexture2D* UHightScreenshotLib::CreateTexture2DFromImage(int sizeX, int sizeY
	, const TArray<FColor>& ColorArray, EColorGammaSpace gammaSpace)
{
	if (sizeX* sizeY!= ColorArray.Num())
	{
		return nullptr;
	}
	FImageView ImageView(ColorArray.GetData(), sizeX, sizeY,EGammaSpace(gammaSpace));
	return FImageUtils::CreateTexture2DFromImage(ImageView);
}
UTexture2D* UHightScreenshotLib::CreateTexture2DFromImageL(int sizeX, int sizeY
	, const TArray<FLinearColor>& ColorArray) {
	FImageView ImageView(ColorArray.GetData(), sizeX, sizeY);
	return FImageUtils::CreateTexture2DFromImage(ImageView);
}